# Bayesian heteroskedastic linear model
# Coded by Andreas Murr, CIDE & University of Warwick, andreas.murr@cide.edu & a.murr@warwick.ac.uk
# Implements algorithm by Cepeda, E. & Gamerman, D. (2000), ‘Bayesian modeling of variance heterogeneity in normal regression models’, Brazilian Journal of Probability and Statistics 14(2), 207–221.

hlm = function(y, X, Z, n.iter = 10000, v = 100, n.prog = 1000){
  
  # load MASS package for mvrnorm()

  library(MASS)

  # density of multivariate normal

  dmvnorm = function(x, mu, Sigma){
    k = length(mu)
    (2 * pi)^(-k/2) * det(Sigma)^(-1/2) * exp(-t(x - mu) %*% solve(Sigma) %*% (x - mu))  
  }

  # set priors

  k = ncol(X)
  p = ncol(Z)
  b.0 = rep(0, k)
  B = v * diag(k)
  B.inv = 1 / v * diag(k)
  g.0 = rep(0, p)
  G = v * diag(p)
  G.inv = 1 / v * diag(p)

  # derive conditional distributinos

  b.g = b.0
  B.g = B
  g.b = g.0
  G.b = G
  B.g.inv = B.inv
  G.b.inv = G.inv

  # compute constants

  n = nrow(Z)
  t.X = t(X)
  t.Z = t(Z)
  B.g.i.b.g = B.g.inv%*%b.g
  G.b.star = solve(G.b.inv + t.Z%*%Z / 2)
  G.b.i.g.b = G.b.inv%*%g.b

  # create vector to save acceptance

  a = rep(NA, n.iter + 1)

  # create vectors to save parameters
  
  beta = matrix(NA, nrow = n.iter + 1, ncol = k)
  gamma = matrix(NA, nrow = n.iter + 1, ncol = p)
  
  # assign starting values to gamma

  ll = function(theta, y, X, Z){
    mu = X%*%theta[1:ncol(X)]
    sigma = sqrt(exp(Z%*%theta[(ncol(X)+1):length(theta)]))
    -sum(dnorm(y, mu, sigma, log = TRUE))
  }
  beta.ini = c(mean(y), rep(0, ncol(X)-1))
  gamma.ini = rep(0, ncol(Z))
  theta.ini = c(beta.ini, gamma.ini)
  fit = optim(theta.ini, ll, hessian = F, method = "BFGS", y = y, X = X, Z = Z)

  # gamma[1,] = 0
  # gamma[1,] = c(log(var(y)), rep(0, p - 1))
  gamma[1,] = fit$par[(ncol(X)+1):length(theta.ini)]
  # gamma[i,] = mvrnorm(1, g.0, G)
  
  # metropolis hastings

  for (i in 1:(n.iter)){

    # beta

    Z.gamma = Z%*%gamma[i,]
    sigma.sq = exp(Z.gamma)
    # Sigma = diag(c(sigma.sq))
    Sigma.inv = 1 / c(sigma.sq) * diag(n)
    t.X.Sigma.inv = t.X%*%Sigma.inv
    B.star = solve(B.g.inv + t.X.Sigma.inv%*%X)
    b.star = B.star %*% (B.g.i.b.g + t.X.Sigma.inv%*%y)
    beta[i,] = mvrnorm(1, b.star, B.star)

    # gamma
    X.beta = X%*%beta[i,]
    y.tilde = Z.gamma + (y - X.beta)^2 / sigma.sq - 1
    g.b.star = G.b.star %*% (G.b.i.g.b + t.Z%*%y.tilde / 2)
    ## propose new value
    
    gamma.prop = mvrnorm(1, g.b.star, G.b.star)

    ## accept?
    
    log.c = sum(dnorm(y, mean = X.beta, sd = sqrt(c(sigma.sq)), log = T)) + log(dmvnorm(c(beta[i,], gamma[i,]), c(g.0, b.0), v * diag(ncol(X) + ncol(Z))))
    log.n = sum(dnorm(y, mean = X.beta, sd = sqrt(exp(Z%*%gamma.prop)), log = T)) + log(dmvnorm(c(beta[i,], gamma.prop), c(g.0, b.0), v * diag(k + p)))
    if(log.n == -Inf){
      a[i] = 0
      gamma[i+1,] = gamma[i,]
    } else{
      if(runif(1) < exp(log.n - log.c)){
        gamma[i+1,] = gamma.prop
        a[i] = 1
      } else{
        gamma[i+1,] = gamma[i,]
        a[i] = 0
      }
    }
    if(i %% n.prog==0) {
       # Print on the screen some message
       cat(paste0("iteration: ", i, "\n"))
    }  
  }
  list(beta[-(n.iter+1),], gamma[-(n.iter+1),], a[-(n.iter+1)])
}

plot.hlm.iter = function(fit, n.burn = 0, store = 1){
  n.iter = length(fit[[3]])
  if (n.burn == 0){
    beta.save = fit[[1]][seq(1, (n.iter-n.burn), store),]
    gamma.save = fit[[2]][seq(1, (n.iter-n.burn), store),]  
  } else{
    beta.save = fit[[1]][-(1:n.burn),][seq(1, (n.iter-n.burn), store),]
    gamma.save = fit[[2]][-(1:n.burn),][seq(1, (n.iter-n.burn), store),]      
  }
  iter = 1:nrow(beta.save)
  K = ncol(beta.save)
  P = ncol(gamma.save)
  par(las = 1, mfcol = c(K, 2), mar = c(1,4,1,1), oma = c(2,0,1,0))
  for (k in 1:K){
    plot(iter, beta.save[,k], main = "", type = "l", xlab = "", ylab = "", xaxt = "n")
    axis(1, labels = F)
    mtext(bquote(beta[.(k)]), 3, 0)
    if(k == K){
      mtext("Iteration", 1, 2) 
      axis(1, labels = T)
    }
  }
  for (p in 1:P){
    plot(iter, gamma.save[,p], main = "", type = "l", xlab = "", ylab = "", xaxt = "n")
    axis(1, labels = F)
    mtext(bquote(gamma[.(p)]), 3, 0)
    if(p == P){
      mtext("Iteration", 1, 2) 
      axis(1, labels = T)
    }
  }
}

summary.hlm = function(fit, n.burn = 1, store = 1, y, X){
  # assumes same predictors in X and Y
  n.iter = length(fit[[3]])
  beta.save = fit[[1]][-(1:n.burn),][seq(1, (n.iter-n.burn), store),]
  gamma.save = fit[[2]][-(1:n.burn),][seq(1, (n.iter-n.burn), store),]
  b.est = apply(beta.save, 2, mean)
  b.std = apply(beta.save, 2, sd)
  b.qua = apply(beta.save, 2, function(x){quantile(x, c(0.025, .975))})
  g.est = apply(gamma.save, 2, mean)
  g.std = apply(gamma.save, 2, sd)
  g.qua = apply(gamma.save, 2, function(x){quantile(x, c(0.025, .975))})
  a = mean(fit[[3]])
  # create output
  out = matrix(NA, nrow = length(b.est) * 3, ncol = 2)
  ## beta
  out[seq(1, nrow(out), 3), 1] = formatC(b.est, format = "f", digits = 2)
  out[seq(2, nrow(out), 3), 1] = paste("(", formatC(b.std, format = "f", digits = 2), ")", sep = "")
  out[seq(3, nrow(out), 3), 1] = paste("[", formatC(b.qua[1,], format = "f", digits = 2), "; ", formatC(b.qua[2,], format = "f", digits = 2), "]", sep = "")
  # gamma
  out[seq(1, nrow(out), 3), 2] = formatC(g.est, format = "f", digits = 2)
  out[seq(2, nrow(out), 3), 2] = paste("(", formatC(g.std, format = "f", digits = 2), ")", sep = "")
  out[seq(3, nrow(out), 3), 2] = paste("[", formatC(g.qua[1,], format = "f", digits = 2), "; ", formatC(g.qua[2,], format = "f", digits = 2), "]", sep = "")
  ## annotation
  colnames(out) = c("Mean", "Variance")
  rownames(out) = rep("", nrow(out))
  rownames(out)[seq(1, nrow(out), 3)] = colnames(X)
  out = rbind(out, c(paste("n = ", nrow(X), sep = ""), paste("ll = ", formatC(sum(dnorm(y, mean = X%*%b.est, sd = sqrt(exp(X%*%g.est)), log = T)), format = "f", digits = 2), sep = ""))) 
  ## return
  noquote(out)
}

# ===================
# = end source code =
# ===================